#pragma once

#include <iostream>
#include <unordered_set>
#include <string>
#include <sys/types.h>
#include <sys/socket.h>
#include <cstring>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>
#include <sys/wait.h>

using std::cout;
using std::cin;
using std::cerr;
using std::endl;

void childFdJoin(int sig)
{
    cout << "回收子进程" << endl;
    while(waitpid(-1, nullptr, 1));  // 回收子进程
}

class TCPServer
{
private:
    int _listenSock;  // 监听套接字
    std::unordered_set<std::string> _ipBackMap;  // ip黑名单
public:
    TCPServer(uint16_t port)
    {
        // 构建监听套接字
        _listenSock = socket(AF_INET, SOCK_STREAM, 0);
        if (_listenSock < 0)
        {
            cerr << "socket error " << strerror(errno) << endl; 
            exit(1);
        }
        // 绑定
        struct sockaddr_in server;
        bzero(&server, sizeof server);
        server.sin_family = AF_INET;
        server.sin_addr.s_addr = inet_addr("0.0.0.0");
        server.sin_port = htons(port);

        if (0 > bind(_listenSock, (sockaddr*)&server, sizeof server))
        {
            cerr << "bind error " << strerror(errno) << endl; 
            exit(2);
        }
        // 设置监听套接字和连接队列长度
        if (0 > listen(_listenSock, 20))
        {
            cerr << "listen error " << strerror(errno) << endl; 
            exit(3);
        }
        
        // 增加功能，加载黑名单IP
        int fd = open("./IpBack.conf", O_RDONLY | O_CREAT, 0664);  // rw-rw-r-- 664
        char ar;
        std::string arr;
        int n = read(fd, &ar, 1);  // "127.0.0.1
        while (n == 1)
        {
            // 读取成功继续
            if (ar != '\n')
            {
                arr += ar;
            }
            else{
                _ipBackMap.insert(arr);
                arr.clear();
            }
            n = read(fd, &ar, 1);
        }
        cout << "Ip黑名单加载完毕" << endl;

        close(fd);
        cout << "TCP服务器初始化成功" << endl;
    }

    void start()
    {
        // 启动
        // 注册信号处理信号使其自动回收
        signal(SIGCHLD, childFdJoin);

        // TCP首先选择一条连接上来
        struct sockaddr_in client;
        bzero(&client, 0);
        socklen_t client_size = sizeof client;

        while (true)
        {
            int sock = accept(_listenSock, (sockaddr*)&client, &client_size);
            if (sock < 0)
            {
                cerr << "提取连接失败 accept error " << strerror(errno) << endl;
                exit(4);
            }
            cout << "连接获取成功\n";
            // 提取成功，判断是不是黑名单的IP
            std::string clientIp = inet_ntoa(client.sin_addr);
            uint16_t clientPort = ntohs(client.sin_port);

            if (_ipBackMap.find(clientIp) != _ipBackMap.end())
            {
                close(sock);  // 服务器主动关闭连接
                continue;
            }

            // 下面创建子进程处理这条连接
            int fd = fork();
            if(fd < 0)
            {
                cerr << "fork error " << strerror(errno) << "\n";
                exit(5);
            }
            else if (fd == 0)
            {
                // 子进程
                // 关闭监听套接字的影响
                close(_listenSock);
                char buffer[1024];
                while (true)
                {
                    ssize_t n = read(sock, buffer, 1024);
                    if (n > 0)
                    {
                        buffer[n] = '\0';
                        cout << "[" << clientIp << ": " << clientPort << "]# " << buffer << "\n"; 
                    }
                    else if (n == 0)
                    {
                        // 对方断开连接
                        cout << "对方断开连接！" << endl;
                        close(sock);
                        break;
                    }
                    else
                    {
                        // 读取失败
                        cerr << "读取失败! read error" << strerror(errno) << endl;
                        close(sock);
                        break;
                    }

                    send(sock, buffer, n, 0);  // 发送回去
                }
                exit(0);
            }

            // 父进程最后别忘了关闭连接
            close(sock);
        }
    }

    ~TCPServer()
    {
        close(_listenSock);
    }
};